import os
import pytest
from unittest.mock import patch
from motor.utils.http_client import SafeHTTPSClient
import requests
import copy
from urllib.parse import urlparse

from motor.resources.instance import Instance, PDRole, Endpoint
from motor.coordinator.core.instance_manager import InstanceManager
from motor.coordinator.metrics.metrics_collector import MetricsCollector, MetricType, SingleMetric


class MockResponse:
    def __init__(self, json_data, status_code):
        self.json_data = json_data
        self.status_code = status_code

    def json(self):
        return self.json_data

class TestMetrics:
    def setup_method(self):
        ep0 = Endpoint(id=0, ip="127.0.0.1", port="8000")
        ep1 = Endpoint(id=1, ip="127.0.0.1", port="8001")
        ep2 = Endpoint(id=2, ip="127.0.0.1", port="8002")
        ep3 = Endpoint(id=3, ip="127.0.0.1", port="8003")
        ep4 = Endpoint(id=4, ip="127.0.0.1", port="8004")
        ep5 = Endpoint(id=5, ip="127.0.0.1", port="8005")
        self.p_ins = Instance(
            job_name="test-prefill",
            model_name="test-model",
            id=0,
            role=PDRole.ROLE_P,
            endpoints={
                "127.0.0.1": { 0: ep0, 1: ep1 }
            }
        )
        self.d_ins = Instance(
            job_name="test-decode",
            model_name="test-model",
            id=1,
            role=PDRole.ROLE_D,
            endpoints={
                "127.0.0.1": { 2: ep2, 3: ep3 }
            }
        )
        self.h_ins = Instance(
            job_name="test-hybrid",
            model_name="test-model",
            id=2,
            role=PDRole.ROLE_U,
            endpoints={
                "127.0.0.1": { 4: ep4, 5: ep5 }
            }
        )

        self.metrics_template = self.load_example_metrics()

    def load_example_metrics(self):
        script_path = os.path.abspath(__file__)
        script_dir = os.path.dirname(script_path)
        data_path = os.path.join(script_dir, "metrics_example.txt")
        with open(data_path, 'r') as f:
            return f.read().strip()

    def clean_instances(self):
        if hasattr(MetricsCollector(), '_initialized'):
            MetricsCollector().stop()
            del MetricsCollector()._initialized

    def load_test_gauge_metric(self):
        # metric text
        metric_str_gauge = """
# HELP vllm:num_requests_running Number of requests in model execution batches.
# TYPE vllm:num_requests_running gauge
vllm:num_requests_running{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 1.0"""

        # metric format
        metric_gauge = SingleMetric()
        metric_gauge.name = "vllm:num_requests_running"
        metric_gauge.help = "Number of requests in model execution batches."
        metric_gauge.type = MetricType.GAUGE
        metric_gauge.label = [
            'vllm:num_requests_running{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}'
        ]
        metric_gauge.value = [1.0]

        return metric_str_gauge.strip(), metric_gauge

    def load_test_counter_metric(self):
        # metric text
        metric_str_counter = """
# HELP vllm:request_success_total Count of successfully processed requests.
# TYPE vllm:request_success_total counter
vllm:request_success_total{engine="0",finished_reason="stop",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 1.0
vllm:request_success_total{engine="0",finished_reason="length",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 2.0
vllm:request_success_total{engine="0",finished_reason="abort",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 0.0"""

        # metric format
        metric_counter = SingleMetric()
        metric_counter.name = "vllm:request_success_total"
        metric_counter.help = "Number of requests in model execution batches."
        metric_counter.type = MetricType.COUNTER
        metric_counter.label = [
            'vllm:request_success_total{engine="0",finished_reason="stop",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_success_total{engine="0",finished_reason="length",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_success_total{engine="0",finished_reason="abort",model_name="/job/model/Qwen2.5-0.5B-Instruct"}'
        ]
        metric_counter.value = [1.0, 2.0, 0.0]

        return metric_str_counter.strip(), metric_counter

    def load_test_histogram_metric(self):
        # metric text
        metric_str_histogram = """
# HELP vllm:request_params_n Histogram of the n request parameter.
# TYPE vllm:request_params_n histogram
vllm:request_params_n_bucket{engine="0",le="1.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_bucket{engine="0",le="2.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_bucket{engine="0",le="5.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_bucket{engine="0",le="10.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_bucket{engine="0",le="20.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_bucket{engine="0",le="+Inf",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_count{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0
vllm:request_params_n_sum{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 3.0"""

        # metric format
        metric_histogram = SingleMetric()
        metric_histogram.name = "vllm:request_params_n"
        metric_histogram.help = "Histogram of the n request parameter."
        metric_histogram.type = MetricType.HISTOGRAM
        metric_histogram.label = [
            'vllm:request_params_n_bucket{engine="0",le="1.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_bucket{engine="0",le="2.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_bucket{engine="0",le="5.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_bucket{engine="0",le="10.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_bucket{engine="0",le="20.0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_bucket{engine="0",le="+Inf",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_count{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}',
            'vllm:request_params_n_sum{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"}'
        ]
        metric_histogram.value = [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0]

        return metric_str_histogram.strip(), metric_histogram

    def load_test_summary_metric(self):
        metric_str_summary = """
# HELP http_request_size_bytes Content length of incoming requests by handler. Only value of header is respected. Otherwise ignored. No percentile calculated.
# TYPE http_request_size_bytes summary
http_request_size_bytes_count{handler="/v1/completions"} 2.0
http_request_size_bytes_sum{handler="/v1/completions"} 312.0
http_request_size_bytes_count{handler="/v1/chat/completions"} 1.0
http_request_size_bytes_sum{handler="/v1/chat/completions"} 268.0"""

        metric_summary = SingleMetric()
        metric_summary.name = "http_request_size_bytes"
        metric_summary.help = "Content length of incoming requests by handler. Only value of header is respected. Otherwise ignored. No percentile calculated."
        metric_summary.type = MetricType.SUMMARY
        metric_summary.label = [
            'http_request_size_bytes_count{handler="/v1/completions"}',
            'http_request_size_bytes_sum{handler="/v1/completions"}',
            'http_request_size_bytes_count{handler="/v1/chat/completions"}',
            'http_request_size_bytes_sum{handler="/v1/chat/completions"}'
        ]
        metric_summary.value = [2.0, 312.0, 1.0, 268.0]

        return metric_str_summary.strip(), metric_summary

    def check_metric_value_equel(self, a, b):
        if not isinstance(a, list) or not isinstance(b, list):
            return False

        if len(a) != len(b):
            return False

        allow_diff = 0.01
        for i in range(len(a)):
            if not isinstance(a[i], float) or not isinstance(b[i], float):
                return False
            if abs(a[i] - b[i]) > allow_diff:
                return False

        return True

    def test_parse_metrics_text_normal(self):
        metric_collector = MetricsCollector()

        # load test metric data
        metric_list = [
            self.load_test_gauge_metric(),
            self.load_test_counter_metric(),
            self.load_test_histogram_metric(),
            self.load_test_summary_metric()
        ]

        # create mix data of 4 type metrics
        merged_metric_str = ""
        merged_metric = []
        for metric_str, metric in metric_list:
            merged_metric_str += metric_str
            merged_metric.append(merged_metric)

        # check _parse_metric_text use metric_list
        for metric_str, metric in metric_list:
            result = metric_collector._parse_metric_text(metric_str)
            assert isinstance(result, list)
            assert len(result) == 1
            assert result[0].name == metric.name
            assert result[0].type == metric.type
            assert result[0].label == metric.label
            assert self.check_metric_value_equel(result[0].value, metric.value)

        # check _parse_metric_text use full metric data
        metric_collector = MetricsCollector()
        result = metric_collector._parse_metric_text(self.metrics_template)
        assert isinstance(result, list)
        assert len(result) > 0

    def test_parse_metrics_text_abnormal(self):
        metrics_str_type_error = """
# HELP vllm:num_requests_running Number of requests in model execution batches.
# TYPE vllm:num_requests_running type_error
vllm:num_requests_running{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} 1.0"""

        metrics_str_value_type_error = """
# HELP vllm:num_requests_running Number of requests in model execution batches.
# TYPE vllm:num_requests_running illegal_type
vllm:num_requests_running{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} value_type_error"""

        metrics_str_value_error = """
# HELP vllm:num_requests_running Number of requests in model execution batches.
# TYPE vllm:num_requests_running illegal_type
vllm:num_requests_running{engine="0",model_name="/job/model/Qwen2.5-0.5B-Instruct"} -1.0"""

        metric_collector = MetricsCollector()
        result = metric_collector._parse_metric_text(metrics_str_type_error)
        assert isinstance(result, list)
        assert len(result) == 0

        result = metric_collector._parse_metric_text(metrics_str_value_type_error)
        assert isinstance(result, list)
        assert len(result) == 0

        result = metric_collector._parse_metric_text(metrics_str_value_error)
        assert isinstance(result, list)
        assert len(result) == 0

    def test_clear_inactive_metrics(self):
        # ensure MetricsCollector clean
        self.clean_instances()
        metric_collector = MetricsCollector()

        # create a gauge type metric
        _, metric_gauge = self.load_test_gauge_metric()
        _, metric_counter = self.load_test_counter_metric()
        _, metric_histogram = self.load_test_histogram_metric()
        _, metric_summary = self.load_test_summary_metric()

        metric_collector._clear_inactive_metrics({})
        assert len(metric_collector._inactive_instance_metrics_aggregate) == 0

        unavailable_pool = {
            self.p_ins.id: self.p_ins
        }
        metric_collector._clear_inactive_metrics(unavailable_pool)
        assert len(metric_collector._inactive_instance_metrics_aggregate) == 0

        metric_collector._instance_metrics_cached = {
            self.p_ins.id: {
                "metrics": [
                    metric_gauge,
                    metric_counter,
                    metric_histogram,
                    metric_summary
                ]
            }
        }
        metric_collector._clear_inactive_metrics(unavailable_pool)
        assert len(metric_collector._instance_metrics_cached) == 0
        assert self.check_metric_value_equel(
                        metric_collector._inactive_instance_metrics_aggregate[0].value, 
                        [0.0] * len(metric_gauge.value)
                    )
        assert self.check_metric_value_equel(
                        metric_collector._inactive_instance_metrics_aggregate[1].value, 
                        metric_counter.value
                    )
        assert self.check_metric_value_equel(
                        metric_collector._inactive_instance_metrics_aggregate[2].value, 
                        metric_histogram.value
                    )
        assert self.check_metric_value_equel(
                        metric_collector._inactive_instance_metrics_aggregate[3].value, 
                        metric_summary.value
                    )

    def mock_get_all_instances_normal(self):
        available_pool = {
            self.p_ins.id: self.p_ins,
            self.d_ins.id: self.d_ins,
            self.h_ins.id: self.h_ins,
        }

        unavailable_pool = {}
        return available_pool, unavailable_pool

    def mock_requests_get_normal(self, *args, **kwargs):
        return MockResponse(self.metrics_template, 200)

    @patch('motor.coordinator.core.instance_manager.InstanceManager.get_all_instances')
    def test_get_all_instances(self, mock_get_all_instances):
        mock_get_all_instances.side_effect = self.mock_get_all_instances_normal

        assert InstanceManager().get_all_instances() == self.mock_get_all_instances_normal()

    @patch('requests.get')
    def test_requests_get(self, mock_requests_get):
        mock_requests_get.side_effect = self.mock_requests_get_normal

        for port in [8000, 8001, 8002, 8003, 8004, 8005]:
            assert requests.get(f"http://localhost:{port}/metrics").status_code == 200

        mock_requests_get.side_effect = self.mock_requests_get_with_abnormal

        for port in [8000, 8001, 8002, 8003]:
            assert requests.get(f"http://localhost:{port}/metrics").status_code == 200
        for port in [8004, 8005]:
            assert requests.get(f"http://localhost:{port}/metrics").status_code == 404

    @patch('motor.coordinator.core.instance_manager.InstanceManager.get_all_instances')
    @patch('requests.get')
    def test_prometheus_metrics_handler(self, mock_requests_get, mock_get_all_instances):
        mock_get_all_instances.side_effect = self.mock_get_all_instances_normal
        mock_requests_get.side_effect = self.mock_requests_get_normal

        self.clean_instances()
        metric_collector = MetricsCollector()

        result = metric_collector.prometheus_metrics_handler()
        assert result is not None

        result = metric_collector.prometheus_instance_metrics_handler()
        assert result is not None

    def mock_requests_get_with_abnormal(self, *args, **kwargs):
        port = urlparse(args[0]).port
        if port in [8000, 8001, 8002, 8003]:
            return MockResponse(self.metrics_template, 200)
        return MockResponse(None, 404)

    @patch('motor.coordinator.core.instance_manager.InstanceManager.get_all_instances')
    @patch('requests.get')
    def test_prometheus_metrics_handler_abnormal(self, mock_requests_get, mock_get_all_instances):
        mock_get_all_instances.side_effect = self.mock_get_all_instances_normal
        mock_requests_get.side_effect = self.mock_requests_get_with_abnormal # abnormal

        self.clean_instances()
        metric_collector = MetricsCollector()

        result = metric_collector.prometheus_metrics_handler()
        assert result is not None

        result = metric_collector.prometheus_instance_metrics_handler()
        assert result is not None
